import os
import cv2
import numpy as np
import json


def resample_image(image, target_resolution):
    """
    将图像重新采样到目标分辨率。

    :param image: 输入图像。
    :param target_resolution: 目标分辨率 (宽度, 高度)。
    :return: 重新采样后的图像。
    """
    return cv2.resize(image, target_resolution, interpolation=cv2.INTER_LINEAR)


def crop_image_random(image, crop_size):
    """
    将图像随机裁剪到指定大小。

    :param image: 输入图像。
    :param crop_size: 裁剪大小 (宽度, 高度)。
    :return: 裁剪后的图像和裁剪位置 (start_x, start_y)。
    """
    h, w = image.shape[:2]
    ch, cw = crop_size
    start_x = np.random.randint(0, w - cw + 1)
    start_y = np.random.randint(0, h - ch + 1)
    cropped_image = image[start_y:start_y + ch, start_x:start_x + cw]
    return cropped_image, (start_x, start_y)


def preprocess_images(reference_image, template_image, target_resolution, crop_size):
    """
    对参考图像和模板图像进行预处理。

    :param reference_image: 参考图像。
    :param template_image: 模板图像。
    :param target_resolution: 目标分辨率 (宽度, 高度)。
    :param crop_size: 模板图像的裁剪大小 (宽度, 高度)。
    :return: 预处理后的参考图像和模板图像以及裁剪位置。
    """
    ref_resampled = resample_image(reference_image, target_resolution)
    template_resampled = resample_image(template_image, target_resolution)
    template_cropped, crop_position = crop_image_random(template_resampled, crop_size)
    return ref_resampled, template_cropped, crop_position


def process_and_save_images(input_dir, output_dir, target_resolution, crop_size):
    """
    处理并保存图像到指定目录。

    :param input_dir: 输入图像目录。
    :param output_dir: 输出图像目录。
    :param target_resolution: 目标分辨率 (宽度, 高度)。
    :param crop_size: 模板图像的裁剪大小 (宽度, 高度)。
    """
    ref_input_dir = os.path.join(input_dir, 'reference')
    temp_input_dir = os.path.join(input_dir, 'template')
    ref_output_dir = os.path.join(output_dir, 'reference')
    temp_output_dir = os.path.join(output_dir, 'template')
    positions_output_file = os.path.join(output_dir, 'crop_positions.json')

    # 检查并创建输出目录
    os.makedirs(ref_output_dir, exist_ok=True)
    os.makedirs(temp_output_dir, exist_ok=True)

    # 检查输入目录是否存在
    if not os.path.exists(ref_input_dir):
        raise FileNotFoundError(f"目录不存在: {ref_input_dir}")
    if not os.path.exists(temp_input_dir):
        raise FileNotFoundError(f"目录不存在: {temp_input_dir}")

    ref_images = os.listdir(ref_input_dir)
    temp_images = os.listdir(temp_input_dir)

    crop_positions = []

    for ref_img_name, temp_img_name in zip(ref_images, temp_images):
        ref_img_path = os.path.join(ref_input_dir, ref_img_name)
        temp_img_path = os.path.join(temp_input_dir, temp_img_name)

        ref_image = cv2.imread(ref_img_path, cv2.IMREAD_COLOR)
        temp_image = cv2.imread(temp_img_path, cv2.IMREAD_COLOR)

        ref_processed, temp_processed, crop_position = preprocess_images(ref_image, temp_image, target_resolution, crop_size)

        cv2.imwrite(os.path.join(ref_output_dir, ref_img_name), ref_processed)
        cv2.imwrite(os.path.join(temp_output_dir, temp_img_name), temp_processed)

        crop_positions.append({
            "filename": temp_img_name,
            "Tx": crop_position[0],
            "Ty": crop_position[1]
        })

    # 保存裁剪位置到JSON文件
    with open(positions_output_file, 'w') as f:
        json.dump(crop_positions, f, indent=4)
    print(f"裁剪位置已保存到：{positions_output_file}")


if __name__ == "__main__":
    input_dir = '../data/optical_to_SAR/raw'
    output_dir = '../data/optical_to_SAR/processed'
    target_resolution = (128, 128)
    crop_size = (64, 64)
    process_and_save_images(input_dir, output_dir, target_resolution, crop_size)
    print("图像预处理完成并已保存到", output_dir)
